package example.fuzzer;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.CookieHandler;
import java.net.CookieManager;
import java.net.CookiePolicy;
import java.net.HttpCookie;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import com.gargoylesoftware.htmlunit.FailingHttpStatusCodeException;
import com.gargoylesoftware.htmlunit.WebClient;
import com.gargoylesoftware.htmlunit.html.HtmlAnchor;
import com.gargoylesoftware.htmlunit.html.HtmlForm;
import com.gargoylesoftware.htmlunit.html.HtmlInput;
import com.gargoylesoftware.htmlunit.html.HtmlPage;
import com.gargoylesoftware.htmlunit.html.HtmlSubmitInput;
import com.gargoylesoftware.htmlunit.html.HtmlTextInput;

/**
 * Authors: Conor Craig, Zach Huenink, Ben Hare
 * 
 * Class is a basic fuzzer that is intended for a specific, small set of web
 * URLs. The Class runs through methods that are turned on or off based on the
 * file conf/config.
 * */
public class BasicFuzzer {
	private static BasicConfig myConfig;

	private static List<HtmlForm> currentHTMLforms_List = null;
	private static ArrayList<String> externalVectors_List = new ArrayList<String>();
	private static ArrayList<String> commonPasswords_List = new ArrayList<String>();
	private static final ArrayList<String> commonlinks = new ArrayList<String>();

	private static List<HtmlAnchor> links;
	private static List<URL> guessedPages;

	// TODO (2pts) Time Gaps.
	// TODO (2pts) Completeness Option (random|full).

	// TODO (5pts) Lack of Sanitization.
	// TODO (5pts) Sensitive Data List.

	// TODO (10pts) Meaningful Output :: Potential Vulnerabilities.
	// TODO (...) Meaningful Output :: Attack Surface.

	private static ArrayList<String> getExternalVectors() {
		return readFromFileIntoArraylist(
				"./conf/" + myConfig.getProperty("vectorsfile"),
				externalVectors_List);
	}

	private static ArrayList<String> getCommonPasswords() {
		ArrayList<String> temp = readFromFileIntoArraylist(
				"./conf/" + myConfig.getProperty("passwordsfile"),
				commonPasswords_List);
		temp.add(0, myConfig.getProperty("password"));
		return temp;
	}

	private static ArrayList<String> readFromFileIntoArraylist(String fileName,
			ArrayList<String> AList) {
		try {
			if (AList.isEmpty()) {
				BufferedReader br = new BufferedReader(new FileReader(new File(
						fileName)));
				if (br.ready()) {
					// System.out.println("Reading from a file...");
					String line;
					while ((line = br.readLine()) != null) {
						AList.add(line);
						// System.out.println("...");
					}
				}
				br.close();
			}
		} catch (IOException e) {
			System.err.println(e.toString());
		}
		return AList;
	}

	/**
	 * The fuzzer should be able to take a URL and parse it down to manipulate
	 * its input and recognize which page a URL goes to. For example,
	 * http://localhost/index.jsp?something=a is the same page as
	 * http://localhost/index.jsp?anotherthing=b, and there are two input that
	 * can be fuzzed (something and anotherthing). Use Java's URL class to help
	 * with this.
	 * 
	 * @param webClient
	 * @throws FailingHttpStatusCodeException
	 * @throws MalformedURLException
	 * @throws IOException
	 */
	private static void parseURLs(String URLtoParse) {
		System.out.println("----------URLs----------");

		String[] URLparts = URLtoParse.split("[*?*]");
		String baseURL = URLparts[0];
		System.out.println("> " + baseURL);

		String[] URLfuzzables = null;
		if (URLparts.length > 1) {
			URLfuzzables = URLparts[1].split("[*&*]");
			for (String s : URLfuzzables) {
				System.out.println("Var> " + s);
			}
		}

		System.out.println("----------EndOfURLs-----------");
	}

	/**
	 * This code finds and prints all the forms of a fuzzed webpage.
	 * 
	 * @param webClient
	 * @throws FailingHttpStatusCodeException
	 * @throws MalformedURLException
	 * @throws IOException
	 */
	private static void discoverFormParameters(WebClient webClient)
			throws FailingHttpStatusCodeException, MalformedURLException,
			IOException {
		if (myConfig.getProperty("discoverforms") == "true") {
			HtmlPage page = webClient.getPage(myConfig.getProperty("url"));
			currentHTMLforms_List = page.getForms();
			System.out.println("------Forms---------");
			for (HtmlForm form : currentHTMLforms_List) {
				System.out.println(form.asText());
			}
			System.out.println("----------EndOfForms-----------");
		}
	}

	/**
	 * *****CUSTOMIZED FOR DVWA ONLY AT THE MOMENT***** Gets username fields and
	 * password fields if there are any and sets them to be admin, and password
	 * respectivly.
	 * 
	 * @param password
	 * @throws IOException
	 */
	private static void authenticateWithPassword(String password)
			throws IOException {
		if (myConfig.getProperty("discoverforms") == "true") {
			HtmlForm actualForm = null;
			for (HtmlForm form : currentHTMLforms_List) {
				actualForm = form
						.getFirstByXPath("//form[@action='login.php']");
				actualForm.setAttribute("username", "admin");
				actualForm.setAttribute("password", password);
				// actualForm.click();
			}
			System.out.println(actualForm.<HtmlPage> click().getWebResponse()
					.getContentAsString());
		}
	}

	/**
	 * 
	 * @param potentialPasswordsList
	 * @throws IOException
	 */
	private static void guessPasswords(List<String> potentialPasswordsList)
			throws IOException {
		if (myConfig.getProperty("passwordguess") != null) {
			for (String password : potentialPasswordsList) {
				authenticateWithPassword(password);
			}
		}
	}

	/**
	 * Page Guessing. The fuzzer should use an external list of common URLs to
	 * discover potentially unlinked pages.
	 * 
	 * @param webpage
	 *            link
	 * */
	// TODO (15 pts) Page & Input Discovery :: Page Guessing.*
	// TODO, create commonLinks
	private static void guessWebPages(WebClient webClient) {
		if (myConfig.getProperty("pageguess") == "true") {
			if (commonlinks != null) {
				for (String commonLink : commonlinks) {
					try {
						String base = getUrlBase(myConfig.getProperty("url"));
						// assemble URL and HTMLPage. (this is the guessing
						// part)
						HtmlPage myPage = webClient.getPage(base + "/"
								+ commonLink);

						// check to see if we already have it.
						// this function will handle adding/discarding
						checkDuplicatePage(myPage);

					} catch (MalformedURLException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					} catch (FailingHttpStatusCodeException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					} catch (IOException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					}
				}
				System.out.println("------Gussed Pages---------");
				for (URL u : guessedPages) {
					System.out.println("Guessed page: " + u.toString());
				}
				System.out.println("------End Guessed Pages---------");

			}
			System.out.println("Detected no common links.");
		}
	}

	/**
	 * Method takes a string URL and returns the string base of that URL.
	 * 
	 * @param String
	 * @return String
	 * */
	private static String getUrlBase(String pagetofuzzString) {
		URL url;
		String base = null;
		try {
			url = new URL(pagetofuzzString);

			// collect indices
			int fileLen = url.getFile().length();
			int strLen = url.toString().length();

			// perform substring operation with indices
			base = url.toString().substring(0, strLen - fileLen);

		} catch (MalformedURLException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		return base;
	}

	/**
	 * Takes an HTMLUnit HtmlPage and checks to see if it already exists in our
	 * discovered links. Add to links if not duplicate, discard if duplicate.
	 * 
	 * @param HtmlPage
	 * */
	private static void checkDuplicatePage(HtmlPage myPage) {
		for (HtmlAnchor h : links) {
			String myLinksURL = h.getHrefAttribute();
			String myParamsURL = myPage.toString();
			if (myLinksURL == myParamsURL) {
				// discard
				return;
			} else {
				try {
					guessedPages.add(new URL(myParamsURL));
				} catch (MalformedURLException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
			}
		}
	}

	/**
	 * Function takes a string URL and returns a set of the cookies associated
	 * with that URL.
	 * 
	 * @param String
	 *            URL
	 * @return Set<Cookie>
	 * */
	private static void discoverCookies(String myURLString) {
		if (myConfig.getProperty("discovercookies") != null) {
			CookieManager cm = new CookieManager();
			cm.setCookiePolicy(CookiePolicy.ACCEPT_ALL);
			CookieHandler.setDefault(cm);

			try {
				new URL(myURLString).openConnection().getContent();
			} catch (MalformedURLException e) {
				e.printStackTrace();
			} catch (IOException e) {
				e.printStackTrace();
			}

			List<HttpCookie> cookies = cm.getCookieStore().getCookies();
			System.out.println("------Cookies---------");
			System.err.println(cookies);
			for (HttpCookie cookie : cookies) {
				System.out.println("Name = " + cookie.getName());
				System.out.println("Value = " + cookie.getValue());
				System.out
						.println("Lifetime (seconds) = " + cookie.getMaxAge());
				System.out.println("Path = " + cookie.getPath());
			}
			System.out.println("------EndOfCookies---------");

		}
	}

	/**
	 * This code is for showing how you can get all the links on a given page,
	 * and visit a given URL
	 * 
	 * @param webClient
	 * @throws IOException
	 * @throws MalformedURLException
	 */
	private static void discoverLinks(WebClient webClient) throws IOException,
			MalformedURLException {
		if (myConfig.getProperty("discoverlinks") == "true") {
			HtmlPage page = webClient.getPage(myConfig.getProperty("url"));

			links = page.getAnchors();
			System.out.println("------Links---------");
			for (HtmlAnchor link : links) {
				System.out.println("Link discovered: " + link.asText()
						+ " @URL=" + link.getHrefAttribute());
			}
			System.out.println("------EndOfLinks---------");
		}
	}

	/**
	 * This code is for demonstrating techniques for submitting an HTML form.
	 * Fuzzer code would need to be more generalized
	 * 
	 * @param webClient
	 * @throws FailingHttpStatusCodeException
	 * @throws MalformedURLException
	 * @throws IOException
	 */
	private static void doFormPost(WebClient webClient)
			throws FailingHttpStatusCodeException, MalformedURLException,
			IOException {
		HtmlPage page = webClient.getPage(myConfig.getProperty("resultsurl"));
		List<HtmlForm> forms = page.getForms();
		for (HtmlForm form : forms) {
			HtmlInput input = form.getInputByName("quantity");
			input.setValueAttribute("2");
			HtmlSubmitInput submit = (HtmlSubmitInput) form
					.getFirstByXPath("//input[@id='submit']");
			System.out.println(submit.<HtmlPage> click().getWebResponse()
					.getContentAsString());
		}
	}

	public static void main(String[] args)
			throws FailingHttpStatusCodeException, MalformedURLException,
			IOException {

		// init config
		File fileIn = new File("./conf/config");
		myConfig = new BasicConfig(fileIn);

		guessedPages = new ArrayList<URL>();
		// TODO create commonLinks
		commonlinks.add("login");
		commonlinks.add("logout");
		commonlinks.add("signup");
		commonlinks.add("index");

		WebClient webClient = new WebClient();
		webClient.setJavaScriptEnabled(true);
		System.err.println();
		parseURLs(myConfig.getProperty("url"));
		getExternalVectors();
		getCommonPasswords();

		discoverFormParameters(webClient);
		authenticateWithPassword(myConfig.getProperty("password"));
		guessPasswords(getCommonPasswords());
		discoverLinks(webClient);
		guessWebPages(webClient);
		discoverCookies(myConfig.getProperty("url"));
		// doFormPost(webClient);

		webClient.closeAllWindows();

	}
}
